# Plot wind intensity time series, 
# Romdhani et al 2022 (Journal of Advances in Modeling Earth Systems)


from matplotlib.ticker import (MultipleLocator, AutoMinorLocator)
import matplotlib.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import csv
import os



#-------------------------------------------------
def Extract_Wind_Speed (Output_File, Wind_Speed_List, Variable_Name):
	with open(Output_File) as f:
		reader = csv.reader(f)
		next (reader)
		row_header = next(reader)

		for row in reader:
			Wind_Speed_List.append(float(row[row_header.index(Variable_Name)]) )
	return (Wind_Speed_List)

#-------------------------------------------------

# This function returns a list of all wrf ouput files in the directory.
def list_csv_files (Dir, csv_files):
	for f in os.listdir(Dir):
		if (f == "Real_Output.csv"):
			csv_files.append(f)
		#if (f == "Dropsonde.csv"):
		#	csv_files.append(f)
		elif (f.find('Smag2D') != -1) and (f.find('cLh1p0') != -1):
			csv_files.append(f)
		elif (f.find('NoTurb') != -1) and (f.find('cLh1p0') != -1):
			csv_files.append(f)
		elif (f.find('TKE2D') != -1) and (f.find('cLh1p0') != -1):
			csv_files.append(f)
	return (csv_files)
#-------------------------------------------------


plt.rcParams.update({
    "text.usetex": True,
    "font.family": "sans-serif",
    "font.sans-serif": ["Helvetica"]})
# for Palatino and other serif fonts use:
plt.rcParams.update({
    "text.usetex": True,
    "font.family": "serif",
    "font.serif": ["Palatino"],
})

Input_Dir = ''
Output_Dir = ''
# Choose between : 'Gustav', 'Irma', 'Katrina', 'Maria'
HNS = ['Katrina']
# Choose the grid resolution
GSS = ['32km']
# Choose the simulation duration.
HR = ['30', '48', '30', '48', '48']
# Choose between : 'NoTurb', 'Smag2D', 'TKE2D'
TMS = ['NoTurb', 'Smag2D', 'TKE2D']
# Choose between : 'cLh0p25', 'cLh0p5', 'cLh1p0', 'cLh1p5'
CLHS = ['cLh1p0']

# List the colors that will be used for tracing the wind intensity.
colors = ['orange', 'red', 'blue', 'black']
# Create a figure
fig, ax = plt.subplots(1, 1, figsize=(10,8), sharex= 'col', sharey= False, )
j = 0
for HN in HNS:
	Times = [0, 6, 12, 18, 24, 30, 36, 42, 48]
	if (HN == 'Gustav'):
		Times = Times [:-3]
	if (HN == 'Katrina'):
		Times = Times [:-3]
		print (Times)
	i = 0
	for GS in GSS:
		# The c incrementation is meant to change the plotting colors.
		c = 0
		csv_files = []
		for CLH in CLHS:
			for TM in TMS:
				Input_Dir_1 = Input_Dir + HN + '/' + GS + '/'
				Hurricane_Setting = HN + '_1Nest_' + HR[j] + 'Hours_MainGrid' + GS + '_' + TM + '_vert42_' + CLH + '_cfl2p0.csv'

			os.chdir(Input_Dir_1)
			csv_files = list_csv_files (Input_Dir_1, csv_files)
			csv_files.sort()
			for csv_file in csv_files:
				print (csv_file)
				Wind_Speed = []
				if (csv_file == "Real_Output.csv"):
					Wind_Speed = Extract_Wind_Speed (csv_file, Wind_Speed, 'Real_Wnd_Ints')
					label = 'Best Track'
				else:
					Wind_Speed = Extract_Wind_Speed (csv_file, Wind_Speed, 'All_Max_WND_SPD_10')
					if (csv_file.find('Smag2D') != -1):
						label = 'Smag2D'
					elif (csv_file.find('TKE2D') != -1):
						label = 'TKE'
					elif (csv_file.find('NoTurb') != -1):
						label = 'NoHorizTurb'


				# Draw segments of of simulated hurricane's wind intensity.
				if label == 'Best Track':
					ax[i, j].plot(Times, Wind_Speed, color=colors[c], linewidth=2, label = label)
				else:
					ax.plot(Times, Wind_Speed, color=colors[c], linewidth=2, linestyle='--', label = label)
				ax.set_yticks(np.arange(20, 90, 10.0))
				ax.set_ylabel("Intensity [m/s]", fontsize=14)
				ax.tick_params(which="minor", axis="x", direction="in")
				ax.tick_params(which="minor", axis="y", direction="in")
				ax.tick_params(which="major", axis="x", direction="in")
				ax.tick_params(which="major", axis="y", direction="in")
				ax.xaxis.set_minor_locator(MultipleLocator(1))
				ax.set_xticks( np.arange(min(Times), max(Times), 6.0))
				ax.yaxis.set_minor_locator(MultipleLocator(1))
				# Define the title
				c = c + 1
		ax.set_title(HN, {'size': 16})
		ax.set_xticks(np.arange(min(Times), max(Times), 6.0))
		ax.set_xlabel("Time [hr]", fontsize=14)
		i = i + 1
	j = j+1

# Create the legend
handles, labels = fig.gca().get_legend_handles_labels()
zipped_lists = zip(labels, handles)
sorted_pairs = sorted(zipped_lists, reverse = True)
tuples = zip(*sorted_pairs)
labels, handles = [ list(tuple) for tuple in  tuples]
by_label = dict(zip(labels, handles))
fig.subplots_adjust(top=0.9)
lgd = fig.legend(by_label.values(), by_label.keys(), bbox_to_anchor=(0.65, 0.99, -0.25, 0), loc = 'upper center', frameon = False, mode='expand', ncol = 4, fontsize='x-large')

# Save the figure
#fig.savefig(Output_Dir + 'Figure_2.eps', format='eps', dpi =600, bbox_extra_artists = (lgd,))
plt.show()